#include <cstdlib>
#include <iostream>
#include <list>
#include <multimedia/imagem/watershed.h>

static const int MASK       = -2;
static const int WSHED      = 0;
static const int INIT       = -1;
static const int FICTITIOUS = -3;

void Init(Pixel***& im, int cx, int cy)
{
    im = (Pixel***) ::malloc(cx*sizeof(Pixel**));
    for(int i=0; i<cx; i++)
    {
        im[i] = (Pixel**) ::malloc(cy*sizeof(Pixel*));
        for(int j=0; j<cy; j++)
        {
            im[i][j] = (Pixel*) ::malloc(sizeof(Pixel));
            im[i][j]->v = 0;
            im[i][j]->x = i;
            im[i][j]->y = j;
        }
    }
}

void Free(Pixel***& im, int cx, int cy)
{
    for(int i=0; i<cx; i++)
    {
        for(int j=0; j<cy; j++)
            ::free(im[i][j]);
        ::free(im[i]);
    }
    ::free(im);
}

void RadixSort(Pixel**& pixels, int w, int h)
{
    Pixel** sorted = (Pixel**) ::malloc(w*h*sizeof(Pixel*));
    
    std::list<Pixel*> lista[256];
    for(int i=0; i<w*h; i++)
        lista[ pixels[i]->v ].push_back(pixels[i]);
    
    int s=0;
    for(int k=0; k<256; k++)
        for(std::list<Pixel*>::iterator it=lista[k].begin(); it!=lista[k].end(); it++)
            sorted[s++] = (*it);
    
    free(pixels);
    pixels = sorted;
}

/*
void Imp(Pixel***& im, int w, int h)
{
    for(int j=0; j<h; j++)
    {
        for(int i=0; i<w; i++)
            printf("%4d ", im[i][j]->v);
        printf("\n");
    }
}
*/

void Watershed_VS(Pixel***& imi, Pixel***& imo, Pixel***& imd, int width, int height)
{
    //
    // Inicialization
    //
    for(int i=0; i<width; i++)
        for(int j=0; j<height; j++)
            imo[i][j]->v = INIT;
    int current_label = 0;
    int current_dist;
    
    for(int i=0; i<width; i++)
        for(int j=0; j<height; j++)
            imd[i][j]->v = 0;
    
    //
    // Sort the pixels of imi in the increasing order of their gray values
    //    
    Pixel** pixels = (Pixel**) ::malloc(width*height*sizeof(Pixel*));
    int s=0;
    for(int j=0; j<height; j++)
        for(int i=0; i<width; i++)
            pixels[s++] = imi[i][j];

    RadixSort(pixels, width, height);
    
    int HMIN = 0;
    int HMAX = 256;
    
    //
    //
    //
    
    std::list<Pixel*> fifo;
    /*Geodesic SKIZ of level h-1 inside level h */
    s=0;
    int t=0;
    for(int h=HMIN; h<HMAX; h++)
    {
        while(s<width*height && pixels[s]->v == h)
        {
            imo[pixels[s]->x][pixels[s]->y]->v = MASK;
            
            for(int m=-1; m<=1; m++)
                for(int n=-1; n<=1; n++)
                {
                    int indiceX = pixels[s]->x + m;
                    int indiceY = pixels[s]->y + n;

                    if(indiceX<0 || indiceX>=width || indiceY<0 || indiceY>=height || (m==0&&n==0))
                        continue;
                    
                    if(imo[indiceX][indiceY]->v > 0 || imo[indiceX][indiceY]->v == WSHED)
                    {
                        imd[pixels[s]->x][pixels[s]->y]->v = 1;
                        fifo.push_back(pixels[s]);
                    }
                }
            
            s++;
        }

        
        current_dist = 1;
    
        // Add FICTITIOUS pixels
        Pixel* fictitious = (Pixel*) ::malloc(sizeof(Pixel));
        fictitious->v = FICTITIOUS;
        fifo.push_back(fictitious);
        
        
        while(true)
        {
            Pixel* p = fifo.front();
            fifo.pop_front();
            
            if(p->v == FICTITIOUS)
            {
                if(fifo.empty())
                    break;
                else
                {
                    fifo.push_back(fictitious);
                    current_dist++;
                    p = fifo.front();
                    fifo.pop_front();
                }
            }
            
            for(int m=-1; m<=1; m++)
                for(int n=-1; n<=1; n++)
                {
                    int indiceX = p->x + m;
                    int indiceY = p->y + n;
    
                    if(indiceX<0 || indiceX>=width || indiceY<0 || indiceY>=height || (m==0&&n==0))
                        continue;
                    
                    if(imd[indiceX][indiceY]->v < current_dist && (imo[indiceX][indiceY]->v>0 || imo[indiceX][indiceY]->v==WSHED))
                    {
                        if(imo[indiceX][indiceY]->v > 0)
                        {
                            if(imo[p->x][p->y]->v == MASK || imo[p->x][p->y]->v == WSHED)
                                imo[p->x][p->y]->v = imo[indiceX][indiceY]->v;
                            else if(imo[p->x][p->y]->v != imo[indiceX][indiceY]->v)
                                imo[p->x][p->y]->v = WSHED;
                        }
                        else if(imo[p->x][p->y]->v == MASK)
                        {
                            imo[p->x][p->y]->v = WSHED;
                        }
                    }
                    else if(imo[indiceX][indiceY]->v == MASK && imd[indiceX][indiceY]->v == 0)
                    {
                        imd[indiceX][indiceY]->v = current_dist + 1;
                        fifo.push_back(imi[indiceX][indiceY]);    ////////////////////
                    }
                }
        }
        
        
        //
        // Checks if a new minimal have been discovered
        //
        while(t<width*height && pixels[t]->v == h)
        {
            imd[ pixels[t]->x ][ pixels[t]->y ]->v = 0;
            
            if(imo[ pixels[t]->x ][ pixels[t]->y ]->v == MASK)
            {
                current_label++;
                fifo.push_back(pixels[t]);
                imo[ pixels[t]->x ][ pixels[t]->y ]->v = current_label;
                
                while(!fifo.empty())
                {
                    Pixel* p_ = fifo.front();
                    fifo.pop_front();
                    
                    for(int m=-1; m<=1; m++)
                        for(int n=-1; n<=1; n++)
                        {
                            int indiceX = p_->x + m;
                            int indiceY = p_->y + n;
            
                            if(indiceX<0 || indiceX>=width || indiceY<0 || indiceY>=height || (m==0&&n==0))
                                continue;
                            
                            if(imo[indiceX][indiceY]->v == MASK)
                            {
                                fifo.push_back(imi[indiceX][indiceY]);
                                imo[indiceX][indiceY]->v = current_label;
                            }
                        }
                }
            }
        
            t++;
        }
    
    }
}
